
import os
import qianfan

# 通过环境变量初始化认证信息
# 方式一：【推荐】使用安全认证AK/SK鉴权
# 替换下列示例中参数，安全认证Access Key替换your_iam_ak，Secret Key替换your_iam_sk，如何获取请查看https://cloud.baidu.com/doc/Reference/s/9jwvz2egb
# os.environ["QIANFAN_ACCESS_KEY"] = os.environ.get("QIANFAN_ACCESS_KEY")
# os.environ["QIANFAN_SECRET_KEY"] = os.environ.get("QIANFAN_SECRET_KEY")

# 方式二：【不推荐】使用应用AK/SK鉴权
# 替换下列示例中参数，将应用API_Key、应用Secret key值替换为真实值
#os.environ["QIANFAN_AK"] = "应用API_Key"
#os.environ["QIANFAN_SK"] = "应用Secret_Key"


class YiYanModel:

    # 通过环境变量初始化认证信息
    # 方式一：【推荐】使用安全认证AK/SK鉴权
    # 替换下列示例中参数，安全认证Access Key替换your_iam_ak，Secret Key替换your_iam_sk，如何获取请查看https://cloud.baidu.com/doc/Reference/s/9jwvz2egb
    # os.environ["QIANFAN_ACCESS_KEY"] = os.environ.get("QIANFAN_ACCESS_KEY")
    # os.environ["QIANFAN_SECRET_KEY"] = os.environ.get("QIANFAN_SECRET_KEY")

    # 方式二：【不推荐】使用应用AK/SK鉴权
    # 替换下列示例中参数，将应用API_Key、应用Secret key值替换为真实值
    #os.environ["QIANFAN_AK"] = "应用API_Key"
    #os.environ["QIANFAN_SK"] = "应用Secret_Key"

    def __init__(self, model='ERNIE-4.0-Turbo-8K'):  # ERNIE-4.0-Turbo-8K, ERNIE-3.5-8K-0613
        self.model = model

    def chat(self, prompt):
        messages = [{'role': 'user', 'content': prompt}]

        chat_comp = qianfan.ChatCompletion()

        # 指定特定模型
        try:
            resp = chat_comp.do(model=self.model, messages=messages)
            return resp["body"]['result']
        except Exception as e:
            print('Error call Yiyan model:', str(e))
            return None


if __name__ == '__main__':
    # call_with_messages()
    llm = YiYanModel()
    prompt_prefix = '按行判断下面每行文字是什么语言。然后将它翻译成英语和简体中文。最后按csv格式输出结果：第一列为原文，第二列为2字母的语言代码，第三列为英语结果，第四列为中文结果。如果原文或译文包含逗号或者引号，按照csv的标准方式进行转义处理。结果只需要输出一遍。不要输出csv表头。如果文字无法翻译，就将这一行舍弃。'
    prompt_content = """
una mente libera
uncinetto che conta
kit gilet jaune et triangle de signalisation voiture
cavo antifurto casco
tete de débroussailleuse husqvarna
portachiavi juventus ufficiale 2024
manette ps5 grise
2 bedroom tiny house
personalised stanley cup
"""
    result = llm.chat(prompt_prefix + '\n' + prompt_content)
    print(result)
